package cn.cosmosx.base.encrypt.sm;

import cn.cosmosx.base.encrypt.hex.HexUtil;
import cn.cosmosx.base.exception.PlatformException;
import lombok.extern.slf4j.Slf4j;
import org.bouncycastle.crypto.CipherParameters;
import org.bouncycastle.crypto.engines.SM4Engine;
import org.bouncycastle.crypto.macs.CBCBlockCipherMac;
import org.bouncycastle.crypto.macs.GMac;
import org.bouncycastle.crypto.modes.GCMBlockCipher;
import org.bouncycastle.crypto.paddings.BlockCipherPadding;
import org.bouncycastle.crypto.paddings.PKCS7Padding;
import org.bouncycastle.crypto.params.KeyParameter;
import org.bouncycastle.crypto.params.ParametersWithIV;
import org.bouncycastle.jce.provider.BouncyCastleProvider;

import javax.crypto.*;
import javax.crypto.spec.IvParameterSpec;
import javax.crypto.spec.SecretKeySpec;
import java.nio.charset.Charset;
import java.nio.charset.StandardCharsets;
import java.security.*;

/**
 * SM4加解密工具类
 *
 * @author gengzhy
 */
@Slf4j
public class SM4Util {
    private final static Charset CHARSET = StandardCharsets.UTF_8;
    /**
     * 默认的SM4秘钥
     */
    private final static String DEFAULT_SM4_SECRET_KEY = "11e22e1b6abe085f71df013eb676d113";

    /**
     * SM4加密
     *
     * @param srcData - 明文
     * @return - 加密后的Hex格式密文
     */
    public static String encryptEBC(String srcData) {
        return encryptEBC(DEFAULT_SM4_SECRET_KEY, srcData);
    }

    /**
     * SM解密
     *
     * @param cipherText - 加密后的Hex密文
     * @return - 明文
     */
    public static String decryptEBC(String cipherText) {
        return decryptEBC(DEFAULT_SM4_SECRET_KEY, cipherText);
    }

    /**
     * @param keyh    16进制文本密钥
     * @param srcData 明文
     * @return 16进制密文
     */
    public static String encryptEBC(String keyh, String srcData) {
        try {
            byte[] key = HexUtil.decode(keyh);
            byte[] cipherText = BaseSM4.encrypt_ECB_Padding(key, srcData.getBytes(StandardCharsets.UTF_8));
            return HexUtil.toHexString(cipherText);
        } catch (Exception ex) {
            ex.printStackTrace();
            log.info("SM4加密失败", ex);
        }
        return null;
    }

    /**
     * ECB模式
     *
     * @param keyh       16进制文本密钥
     * @param cipherText 16进制文本密文
     * @return 明文
     */
    public static String decryptEBC(String keyh, String cipherText) {
        try {
            byte[] key = HexUtil.decode(keyh);
            byte[] decryptedData = BaseSM4.decrypt_ECB_Padding(key, HexUtil.decode(cipherText));
            return new String(decryptedData);
        } catch (Exception ex) {
            log.error("SM4解密失败：", ex);
        }
        return null;
    }

    /**
     * 生成SM4秘钥（SM4算法目前只支持128位，即密钥16个字节）
     *
     * @return - Hex类型的秘钥
     */
    public static String generateKey() {
        try {
            String sm4SecurityKeyHex = HexUtil.toHexString(BaseSM4.generateKey());
            log.info("生成SM4秘钥（HexUtil）：{}", sm4SecurityKeyHex);
            return sm4SecurityKeyHex;
        } catch (Exception ex) {
            log.error("生成SM4秘钥失败", ex);
            throw new PlatformException("生成SM4秘钥失败");
        }
    }

    /*====================================================================*/

    /**
     * @param keyh    16进制文本密钥
     * @param ivh     16进制偏移量
     * @param srcData 明文
     * @return 16进制文本密文
     */
    public static String encryptCBC(String keyh, String ivh, String srcData) {
        try {
            byte[] key = HexUtil.decode(keyh);
            byte[] iv = HexUtil.decode(ivh);
            byte[] cipherText = BaseSM4.encrypt_CBC_Padding(key, iv, srcData.getBytes(CHARSET));
            return HexUtil.toHexString(cipherText);
        } catch (Exception ex) {
            ex.printStackTrace();
            log.info("SM4加密失败：" + ex.getMessage());
        }
        return null;
    }

    /**
     * @param keyh       16进制文本密钥
     * @param ivh        16进制偏移量
     * @param cipherText 16进制文本密文
     * @return 明文
     */
    public static String decryptCBC(String keyh, String ivh, String cipherText) {
        try {
            byte[] key = HexUtil.decode(keyh);
            byte[] iv = HexUtil.decode(ivh);
            byte[] decryptedData = BaseSM4.decrypt_CBC_Padding(key, iv, HexUtil.decode(cipherText));
            return new String(decryptedData);
        } catch (Exception ex) {
            ex.printStackTrace();
            log.info("SM4解密失败：" + ex.getMessage());
        }
        return null;
    }

    /**
     * SM4加解密基础类
     */
    private static class BaseSM4 {
        private static final String ALGORITHM_NAME = "SM4";
        private static final String ALGORITHM_NAME_ECB_PADDING = "SM4/ECB/PKCS5Padding";
        private static final String ALGORITHM_NAME_ECB_NOPADDING = "SM4/ECB/NoPadding";
        private static final String ALGORITHM_NAME_CBC_PADDING = "SM4/CBC/PKCS7Padding";
        private static final String ALGORITHM_NAME_CBC_NOPADDING = "SM4/CBC/NoPadding";
        private static final BouncyCastleProvider PROVIDER = Provider.INSTANCE.getProvider();

        /**
         * SM4算法目前只支持128位（即密钥16字节）
         */
        public static final int DEFAULT_KEY_SIZE = 128;

        public static byte[] generateKey() throws NoSuchAlgorithmException, NoSuchProviderException {
            return generateKey(DEFAULT_KEY_SIZE);
        }

        public static byte[] generateKey(int keySize) throws NoSuchAlgorithmException, NoSuchProviderException {
            KeyGenerator kg = KeyGenerator.getInstance(ALGORITHM_NAME, PROVIDER);
            kg.init(keySize, new SecureRandom());
            return kg.generateKey().getEncoded();
        }

        public static byte[] encrypt_ECB_Padding(byte[] key, byte[] data) throws InvalidKeyException, NoSuchAlgorithmException, NoSuchProviderException, NoSuchPaddingException, IllegalBlockSizeException, BadPaddingException {
            Cipher cipher = generateECBCipher(ALGORITHM_NAME_ECB_PADDING, Cipher.ENCRYPT_MODE, key);
            return cipher.doFinal(data);
        }

        public static byte[] decrypt_ECB_Padding(byte[] key, byte[] cipherText) throws IllegalBlockSizeException, BadPaddingException, InvalidKeyException, NoSuchAlgorithmException, NoSuchProviderException, NoSuchPaddingException {
            Cipher cipher = generateECBCipher(ALGORITHM_NAME_ECB_PADDING, Cipher.DECRYPT_MODE, key);
            return cipher.doFinal(cipherText);
        }

        public static byte[] encrypt_ECB_NoPadding(byte[] key, byte[] data) throws InvalidKeyException, NoSuchAlgorithmException, NoSuchProviderException, NoSuchPaddingException, IllegalBlockSizeException, BadPaddingException {
            Cipher cipher = generateECBCipher(ALGORITHM_NAME_ECB_NOPADDING, Cipher.ENCRYPT_MODE, key);
            return cipher.doFinal(data);
        }

        public static byte[] decrypt_ECB_NoPadding(byte[] key, byte[] cipherText) throws IllegalBlockSizeException, BadPaddingException, InvalidKeyException, NoSuchAlgorithmException, NoSuchProviderException, NoSuchPaddingException {
            Cipher cipher = generateECBCipher(ALGORITHM_NAME_ECB_NOPADDING, Cipher.DECRYPT_MODE, key);
            return cipher.doFinal(cipherText);
        }

        public static byte[] encrypt_CBC_Padding(byte[] key, byte[] iv, byte[] data) throws InvalidKeyException, NoSuchAlgorithmException, NoSuchProviderException, NoSuchPaddingException, IllegalBlockSizeException, BadPaddingException, InvalidAlgorithmParameterException {
            Cipher cipher = generateCBCCipher(ALGORITHM_NAME_CBC_PADDING, Cipher.ENCRYPT_MODE, key, iv);
            return cipher.doFinal(data);
        }

        public static byte[] decrypt_CBC_Padding(byte[] key, byte[] iv, byte[] cipherText) throws IllegalBlockSizeException, BadPaddingException, InvalidKeyException, NoSuchAlgorithmException, NoSuchProviderException, NoSuchPaddingException, InvalidAlgorithmParameterException {
            Cipher cipher = generateCBCCipher(ALGORITHM_NAME_CBC_PADDING, Cipher.DECRYPT_MODE, key, iv);
            return cipher.doFinal(cipherText);
        }

        public static byte[] encrypt_CBC_NoPadding(byte[] key, byte[] iv, byte[] data) throws InvalidKeyException, NoSuchAlgorithmException, NoSuchProviderException, NoSuchPaddingException, IllegalBlockSizeException, BadPaddingException, InvalidAlgorithmParameterException {
            Cipher cipher = generateCBCCipher(ALGORITHM_NAME_CBC_NOPADDING, Cipher.ENCRYPT_MODE, key, iv);
            return cipher.doFinal(data);
        }

        public static byte[] decrypt_CBC_NoPadding(byte[] key, byte[] iv, byte[] cipherText) throws IllegalBlockSizeException, BadPaddingException, InvalidKeyException, NoSuchAlgorithmException, NoSuchProviderException, NoSuchPaddingException, InvalidAlgorithmParameterException {
            Cipher cipher = generateCBCCipher(ALGORITHM_NAME_CBC_NOPADDING, Cipher.DECRYPT_MODE, key, iv);
            return cipher.doFinal(cipherText);
        }

        public static byte[] doCMac(byte[] key, byte[] data) throws NoSuchProviderException, NoSuchAlgorithmException, InvalidKeyException {
            Key keyObj = new SecretKeySpec(key, ALGORITHM_NAME);
            return doMac("SM4-CMAC", keyObj, data);
        }

        public static byte[] doGMac(byte[] key, byte[] iv, int tagLength, byte[] data) {
            org.bouncycastle.crypto.Mac mac = new GMac(new GCMBlockCipher(new SM4Engine()), tagLength * 8);
            return doMac(mac, key, iv, data);
        }

        /**
         * 默认使用PKCS7Padding/PKCS5Padding填充的CBCMAC
         *
         * @param key
         * @param iv
         * @param data
         * @return
         */
        public static byte[] doCBCMac(byte[] key, byte[] iv, byte[] data) {
            SM4Engine engine = new SM4Engine();
            org.bouncycastle.crypto.Mac mac = new CBCBlockCipherMac(engine, engine.getBlockSize() * 8, new PKCS7Padding());
            return doMac(mac, key, iv, data);
        }

        /**
         * @param key
         * @param iv
         * @param padding 可以传null，传null表示NoPadding，由调用方保证数据必须是BlockSize的整数倍
         * @param data
         * @return
         * @throws Exception
         */
        public static byte[] doCBCMac(byte[] key, byte[] iv, BlockCipherPadding padding, byte[] data) throws Exception {
            SM4Engine engine = new SM4Engine();
            if (padding == null) {
                if (data.length % engine.getBlockSize() != 0) {
                    throw new Exception("if no padding, data length must be multiple of SM4 BlockSize");
                }
            }
            org.bouncycastle.crypto.Mac mac = new CBCBlockCipherMac(engine, engine.getBlockSize() * 8, padding);
            return doMac(mac, key, iv, data);
        }

        private static byte[] doMac(org.bouncycastle.crypto.Mac mac, byte[] key, byte[] iv, byte[] data) {
            CipherParameters cipherParameters = new KeyParameter(key);
            mac.init(new ParametersWithIV(cipherParameters, iv));
            mac.update(data, 0, data.length);
            byte[] result = new byte[mac.getMacSize()];
            mac.doFinal(result, 0);
            return result;
        }

        private static byte[] doMac(String algorithmName, Key key, byte[] data) throws NoSuchProviderException, NoSuchAlgorithmException, InvalidKeyException {
            Mac mac = Mac.getInstance(algorithmName, PROVIDER);
            mac.init(key);
            mac.update(data);
            return mac.doFinal();
        }

        private static Cipher generateECBCipher(String algorithmName, int mode, byte[] key) throws NoSuchAlgorithmException, NoSuchProviderException, NoSuchPaddingException, InvalidKeyException {
            Cipher cipher = Cipher.getInstance(algorithmName, PROVIDER);
            Key sm4Key = new SecretKeySpec(key, ALGORITHM_NAME);
            cipher.init(mode, sm4Key);
            return cipher;
        }

        private static Cipher generateCBCCipher(String algorithmName, int mode, byte[] key, byte[] iv) throws InvalidKeyException, InvalidAlgorithmParameterException, NoSuchAlgorithmException, NoSuchProviderException, NoSuchPaddingException {
            Cipher cipher = Cipher.getInstance(algorithmName, PROVIDER);
            Key sm4Key = new SecretKeySpec(key, ALGORITHM_NAME);
            IvParameterSpec ivParameterSpec = new IvParameterSpec(iv);
            cipher.init(mode, sm4Key, ivParameterSpec);
            return cipher;
        }
    }
}
